#!/usr/bin/python -i

import sys
import numpy as np
import matplotlib.pyplot as plt
def HarmonicForce(k,delta):
  return -k*delta
total_time = 100000
time = np.array(range(total_time))
xt  = np.zeros(total_time)
vxt = np.zeros(total_time)
fxt = np.zeros(total_time)

k = 1.0
sigma  = 1.0
xt[0]  = 2.0
vxt[0] = 0.0
fxt[0] = HarmonicForce(k,xt[0])
dt = 0.001
KT = 1.0
etha = np.zeros(total_time)
gaussNum = np.random.normal(0.0, 1, total_time)
gamma = 0.1
for i in range(1,total_time):
  etha[i-1] = np.sqrt(2*KT*dt)*gaussNum[i-1]
  xt[i]  = xt[i-1] + vxt[i-1]*dt + 0.5*fxt[i-1]*dt**2 #+ etha[i-1]
  fxt[i] = HarmonicForce(k,xt[i]) - gamma*vxt[i-1]
  vxt[i] = vxt[i-1] + 0.5*dt*(fxt[i]+fxt[i-1])
  
plt.plot(time,xt)
plt.show()
plt.plot(time,vxt)
plt.show()
xt.mean()
vxt.mean()

for i in range(1,total_time):
  etha[i-1] = np.sqrt(2*KT*dt)*gaussNum[i-1]
  xt[i]  = xt[i-1] + vxt[i-1]*dt + 0.5*fxt[i-1]*dt**2 + etha[i-1]
  fxt[i] = HarmonicForce(k,xt[i]) - gamma*vxt[i-1]
  vxt[i] = vxt[i-1] + 0.5*dt*(fxt[i]+fxt[i-1])
  
plt.plot(time,xt)
plt.show()
plt.plot(time,vxt)
plt.show()
xt.mean()
vxt.mean()
